import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as pch
import time
from scipy.constants import physical_constants

unite_length = 3.4e-10
unite_energy = 1.65e-21
unite_mass = 6.69e-26
unite_time = 2.17e-12
unite_temperature = 120
unite_velocity = 1.57e2

rMinimumSquared = {
    "Lennard-Jones": np.power(2, 1 / 3),
    "Rigid Balls": 1
}
ignoreDistance_2 = 2
k_B = physical_constants["Boltzmann constant"][0]

minimum_simulate_time = 20

SAVE_DIR = os.path.join(os.getcwd(), time.strftime("%Y_%m%d_%H_%M_%S"))
if not os.path.exists(SAVE_DIR): os.makedirs(SAVE_DIR)

def pbcSeparation(difference, L):
    return (difference % -L if difference > 0 else difference % L) if abs(difference) > 0.5 * L else difference

def pbcPosition(position, L):
    return position % L

class LJParticlesApp:

    def __init__(self, Lx=10, Ly=10, Nx=8, Ny=8, N=64, init_T=1, delta_T=0.0001, delta_t=0.01, init_type="rectangle", model="Lennard-Jones"):
        self.Lx, self.Ly, self.Nx, self.Ny, self.init_T, self.delta_T, self.delta_t, self.init_type, self.model = Lx * unite_length, Ly * unite_length, Nx, Ny, init_T * unite_temperature, delta_T * unite_temperature, delta_t * unite_time, init_type, model
        self.N = self.Nx * self.Ny if init_type == "rectangle" else N
        self.record(Lx=Lx, Ly=Ly, Nx=Nx, Ny=Ny, N=self.N, init_T=init_T, delta_T=delta_T, delta_t=delta_t, init_type=init_type, model=model)
        self.delta_t_2 = np.square(self.delta_t)

        self.position_x_ticks = np.linspace(0, self.Lx, self.Nx + 1)
        self.position_y_ticks = np.linspace(0, self.Ly, self.Ny + 1)
        self.position_x_ticks_label = ["{:.2f}".format(_) for _ in np.linspace(0, Lx, self.Nx + 1)]
        self.position_y_ticks_label = ["{:.2f}".format(_) for _ in np.linspace(0, Ly, self.Ny + 1)]

        self.velocity_x = None
        self.velocity_y = None
        self.position_x = None
        self.position_y = None

        self.eq_T = None
        self.simulate_time = None
        self.collision = None

        plt.figure(figsize=(12, 12))
        
    def initialize(self):

        def initialize_velocity():
            random_v = np.sqrt(np.random.rand(64))
            random_phi = (np.random.rand(64) - 0.5) * 2 * np.pi
            self.velocity_x = random_v * np.cos(random_phi)
            self.velocity_y = random_v * np.sin(random_phi)
            vxcm = np.average(self.velocity_x)
            vycm = np.average(self.velocity_y)
            self.velocity_x -= vxcm
            self.velocity_y -= vycm
            kineticEnergy = np.average(np.square(self.velocity_x) + np.square(self.velocity_y)) / 2 * unite_mass
            rescale = np.sqrt(self.init_T / kineticEnergy * k_B)
            self.velocity_x *= rescale
            self.velocity_y *= rescale

        def initialize_position():
            if self.init_type == "random":
                self.position_x = []
                self.position_y = []
                for i in range(self.N):
                    overlap = True
                    position_x = None
                    position_y = None
                    while overlap:
                        overlap = False
                        position_x = self.Lx * np.random.randn()
                        position_y = self.Ly * np.random.randn()
                        for j in range(i):
                            if overlap: break
                            elif np.square(pbcSeparation(position_x - self.position_x[j], self.Lx)) + np.square(pbcSeparation(position_y - self.position_y[j], self.Ly)) <= rMinimumSquared[self.model]:
                                overlap = True
                    self.position_x.append(position_x)
                    self.position_y.append(position_y)
            elif self.init_type == "rectangle":
                self.position_x, self.position_y = np.meshgrid(np.linspace(0, self.Lx, self.Nx, endpoint=False), np.linspace(0, self.Ly, self.Ny, endpoint=False))
                self.position_x = self.position_x.flatten()
                self.position_y = self.position_y.flatten()
        
        np.random.seed(int(time.time()))
        initialize_velocity()
        initialize_position()

    def simulate(self, name=None):
        
        def step():

            def computeAcceleration():
                acceleration_x = np.zeros(self.N)
                acceleration_y = np.zeros(self.N)

                for i in range(self.N):
                    for j in range(i):
                        difference_x = pbcSeparation(self.position_x[i] - self.position_x[j], self.Lx)
                        difference_y = pbcSeparation(self.position_y[i] - self.position_y[j], self.Ly)
                        distance_2 = np.square(difference_x) + np.square(difference_y)
                        distance_2 /= unite_length * unite_length
                        acceleration = unite_energy / (unite_mass * unite_length * unite_length) * (48 / np.power(distance_2, 7) - 24 / np.power(distance_2, 4))

                        acceleration_x_ = difference_x * acceleration
                        acceleration_x_ = 0 if abs(acceleration_x_) > 5e13 else acceleration_x_
                        acceleration_y_ = difference_y * acceleration
                        acceleration_y_ = 0 if abs(acceleration_y_) > 5e13 else acceleration_y_

                        acceleration_x[i] += acceleration_x_
                        acceleration_y[i] += acceleration_y_
                        acceleration_x[j] -= acceleration_x_
                        acceleration_y[j] -= acceleration_y_

                return acceleration_x, acceleration_y

            if self.model == "Lennard-Jones":
                acceleration_x_old, acceleration_y_old = computeAcceleration()
                self.position_x = pbcPosition(self.position_x + self.velocity_x * self.delta_t + acceleration_x_old * self.delta_t_2 / 2, self.Lx)
                self.position_y = pbcPosition(self.position_y + self.velocity_y * self.delta_t + acceleration_y_old * self.delta_t_2 / 2, self.Ly)
                acceleration_x_new, acceleration_y_new = computeAcceleration()
                self.velocity_x += (acceleration_x_new + acceleration_x_old) * self.delta_t / 2
                self.velocity_y += (acceleration_y_new + acceleration_y_old) * self.delta_t / 2
            elif self.model == "Rigid Balls":
                for i in range(self.N):
                    for j in range(i):
                        difference_x = pbcSeparation(self.position_x[i] - self.position_x[j], self.Lx)
                        difference_y = pbcSeparation(self.position_y[i] - self.position_y[j], self.Ly)
                        distance_2 = np.square(difference_x) + np.square(difference_y)
                        distance_2 /= unite_length * unite_length
                        if distance_2 < rMinimumSquared[self.model]:
                            self.collision += 1
                            
                            theta = np.arctan2(difference_y, difference_x)

                            theta_i = np.arctan2(self.velocity_y[i], self.velocity_x[i])
                            difference_theta_i = (theta_i - (theta - np.pi / 2)) % (2 * np.pi)
                            if difference_theta_i > np.pi:
                                theta_i = 2 * np.pi - difference_theta_i + theta - np.pi / 2

                            theta_j = np.arctan2(self.velocity_y[j], self.velocity_x[j])
                            difference_theta_j = (theta_j - (theta - np.pi / 2)) % (2 * np.pi)
                            if difference_theta_j < np.pi:
                                theta_j = 2 * np.pi - difference_theta_j + theta - np.pi / 2

                            velocity_i = np.sqrt(np.square(self.velocity_x[i]) + np.square(self.velocity_y[i]))
                            velocity_j = np.sqrt(np.square(self.velocity_x[j]) + np.square(self.velocity_y[j]))
                            
                            self.velocity_x[i] = velocity_i * np.cos(theta_i)
                            self.velocity_y[i] = velocity_i * np.sin(theta_i)
                            self.velocity_x[j] = velocity_j * np.cos(theta_j)
                            self.velocity_y[j] = velocity_j * np.sin(theta_j)

                self.position_x = pbcPosition(self.position_x + self.velocity_x * self.delta_t, self.Lx)
                self.position_y = pbcPosition(self.position_y + self.velocity_y * self.delta_t, self.Ly)

            if not name is None: self.draw("{}-Step{:d}".format(name, self.simulate_time))
            return np.average(np.square(self.velocity_x) + np.square(self.velocity_y)) / 2 * unite_mass

        self.simulate_time = 1
        self.collision = 0
        E_old = self.init_T * k_B
        delta_E = self.delta_T * k_B
        E_new = step()

        while abs(E_old - E_new) > delta_E or self.simulate_time < minimum_simulate_time:
            E_old = E_new
            E_new = step()
            self.simulate_time += 1

        self.eq_T = E_new / k_B

    def find(self, eq_T, bias=0, limit=None):
        print("Start finding {}K, bias={}".format(eq_T, bias))
        find_time = 0
        while abs(self.eq_T - eq_T) > self.delta_T:
            if not limit is None:
                if find_time > limit: return -1
            self.init_T *= eq_T / self.eq_T + bias
            self.initialize()
            self.simulate()
            find_time += 1
            print(self.eq_T)
        return find_time

    def draw(self, name):
        print("Drawing {}".format(name))
        plt.clf()

        ax1 = plt.subplot(2, 2, 1)
        velocity_x = self.velocity_x / np.max(self.velocity_x) * unite_length * 0.5
        velocity_y = self.velocity_y / np.max(self.velocity_y) * unite_length * 0.5
        for i in range(self.N):
            ax1.arrow(self.position_x[i], self.position_y[i], velocity_x[i], velocity_y[i], width=unite_length * 0.05)
        ax1.scatter(self.position_x, self.position_y)
        rect = pch.Rectangle((0, 0), self.Lx, self.Ly, color="orange", zorder=0)
        ax1.add_patch(rect)
        ax1.set_title("Position and Velocity")
        ax1.set_xticks(self.position_x_ticks)
        ax1.set_xticklabels(self.position_x_ticks_label)
        ax1.set_xlabel("$x/\sigma$")
        ax1.set_yticks(self.position_y_ticks)
        ax1.set_yticklabels(self.position_y_ticks_label)
        ax1.set_ylabel("$y/\sigma$")
        ax1.set_aspect('equal', adjustable='box')

        ax2 = plt.subplot(2, 2, 2)
        ax2.scatter(self.velocity_x, self.velocity_y)
        velocity_max = max(np.max(np.abs(self.velocity_x)), np.max(np.abs(self.velocity_y)))
        N = int(np.ceil(velocity_max / unite_velocity * 2))
        velocity_ticks = np.linspace(-N * unite_velocity / 2, N * unite_velocity / 2, 2 * N + 1)
        velocity_ticks_label = ["{:.1f}".format(_) for _ in np.linspace(-N / 2, N / 2, 2 * N + 1)]
        ax2.set_title("Velocity")
        ax2.set_xticks(velocity_ticks)
        ax2.set_xticklabels(velocity_ticks_label)
        ax2.set_xlabel("$v_{x}/({\epsilon/m}^{1/2})$")
        ax2.set_yticks(velocity_ticks)
        ax2.set_yticklabels(velocity_ticks_label)
        ax2.set_ylabel("$v_{y}/({\epsilon/m}^{1/2})$")
        ax2.set_aspect('equal', adjustable='box')

        ax3 = plt.subplot(2, 1, 2)
        rate = np.sqrt(np.square(self.velocity_x) + np.square(self.velocity_y))
        rate_max = max(rate)
        N = int(np.ceil(rate_max / unite_velocity * 4))
        rate_bin = np.linspace(0, N * unite_velocity / 4, N + 1)
        rate_ticks_label = ["{:.1f}".format(_) for _ in np.linspace(0, N / 4, N + 1)]
        ax3.set_title("Rate")
        ax3.set_xticks(rate_bin)
        ax3.set_xticklabels(rate_ticks_label)
        ax3.set_xlabel("$v/({\epsilon/m}^{1/2})$")
        ax3.set_ylabel("$n$")
        ax3.set_aspect('auto', adjustable='box')
        ax3.hist(rate, bins=rate_bin)

        plt.savefig(os.path.join(SAVE_DIR, "{}.png".format(name)))

    def record(self, **kwargs):
        with open(os.path.join(SAVE_DIR, "result.txt"), 'a') as f:
            string = kwargs.pop("string", None)
            if not string is None: f.write("{}\n".format(string))
            for key, value in kwargs.items():
                f.write("{}: {}\n".format(key, str(value)))

if __name__ == "__main__":
    myapp = LJParticlesApp()
    init_T = myapp.init_T
    delta_T = myapp.delta_T
    myapp.initialize()
    myapp.simulate(name="answer_1_1")
    myapp.record(answer_1_1=myapp.eq_T)
    myapp.delta_T = 0.5
    find_time = myapp.find(400)
    myapp.draw("answer_1_2")
    myapp.record(answer_1_2=myapp.init_T)
    myapp.record(string="answer_2: {:f} {:d}".format(0, find_time))
    for bias in [-1e-1, -1e-2, -1e-3, 1e-3, 1e-2, 1e-1]:
        myapp.init_T = myapp.eq_T = init_T
        find_time = myapp.find(400, bias=bias, limit=50)
        myapp.record(string="answer_2: {:f} {:d}".format(bias, find_time))
    
    myapp.init_T = myapp.eq_T = init_T
    myapp.delta_T = delta_T
    myapp.model = "Rigid Balls"
    myapp.initialize()
    myapp.simulate(name="answer_1_3")
    velocity = np.sum(np.sqrt(np.square(myapp.velocity_x) + np.square(myapp.velocity_y)))
    myapp.record(answer_1_3=velocity * myapp.simulate_time * myapp.delta_t / (myapp.collision * unite_length))